#ifndef SWESOLVER_HPP
#define SWESOLVER_HPP

#include "sweTraits.hpp"
#include "sweSettings.hpp"
#include "solverBase.hpp"

namespace hms
{

/* forward declaration for friend class used for blockwise execution */
namespace swe { template<typename MeshType> class SolverThread; }

/** @brief Solver class for the Shallow Water Equations. */
template<typename MeshType>
class SWESolver : public SolverBase<SWESolver<MeshType>> {
public:
	using Type = SWESolver<MeshType>;
	using BaseType = SolverBase<Type>;
	using typename BaseType::TraitsType;
	using typename BaseType::SettingsType;
	using typename BaseType::ConstantName;
	using typename BaseType::FieldName;
	using typename BaseType::ConstantFieldName;
	using typename BaseType::SourceFieldName;
	using typename BaseType::FluxScheme;
	// using TraitsType = SolverTraits<Type>;

	// using SettingsType      = typename TraitsType::SettingsType;
	// using ConstantName      = typename TraitsType::ConstantName;
	// using FieldName         = typename TraitsType::FieldName;
	// using ConstantFieldName = typename TraitsType::ConstantFieldName;
	// using SourceFieldName   = typename TraitsType::SourceFieldName;
	// using FluxScheme        = typename SettingsType::FluxScheme;
	
	/** @brief The minimum water depth below which a cell is assumed dry. */
	static constexpr scalar minWaterDepth { 1e-8 };

	static constexpr int nConstants      { BaseType::nConstants };
	static constexpr int nFields         { BaseType::nFields };
	static constexpr int nConstantFields { BaseType::nConstantFields };
	static constexpr int nSourceFields   { BaseType::nSourceFields };

	static constexpr bool
		isUniMesh    { std::is_same_v<MeshType, UniMesh> },
		isRectMesh   { std::is_same_v<MeshType, RectMesh> },
		isStructMesh { std::is_same_v<MeshType, StructMesh> },
		isOrthogonal { isUniMesh || isRectMesh },
		isStructured { isOrthogonal || isStructMesh };

private:
	#ifdef SWE_STORE_WATER_ELEVATION
	Field<MeshType> m_waterElevation;
	#endif
	
	#ifdef SOLVER_TIME_FLUXES
	std::array<double, 8> m_execTimesFlux {0, 0, 0, 0, 0, 0, 0, 0};
	#endif

public:

	SWESolver();

	SWESolver( const MeshBase<MeshType>& );

	SWESolver( const MeshBase<MeshType>&, const Time&, scalar uniformManningsN = 0.025 );

	auto rho() const -> scalar;
	void rho( scalar );

	auto nu() const -> scalar;
	void nu( scalar );

	auto v() -> Field<MeshType>&;
	auto v() const -> const Field<MeshType>&;

	auto d() -> Field<MeshType>&;
	auto d() const -> const Field<MeshType>&;

	#ifdef SWE_STORE_WATER_ELEVATION
	auto h() -> Field<MeshType>&;
	auto h() const -> const Field<MeshType>&;
	#endif

	auto q() -> Field<MeshType>&;

	auto z() -> Field<MeshType>&;
	auto z() const -> const Field<MeshType>&;

	auto n() -> ConstantFieldVar<MeshType>&;
	auto n() const -> const ConstantFieldVar<MeshType>&;

	auto r() -> SourceFieldVar<MeshType>&;

	auto i() -> SourceFieldVar<MeshType>&;

	auto isDry() const -> decltype(auto);
	auto isDryDomain() const -> decltype(auto);

	auto useInfiltration() const -> bool;
	void useInfiltration( bool b );

	auto useRainfall() const -> bool;
	void useRainfall( bool b );

	#ifdef SOLVER_TIME_FLUXES
	auto execTimesFlux() -> std::array<double, 8>&;
	#endif

	void updateBoundary( Field<MeshType>&, unsigned patchIndex );

private:
	void setConstants();

	void resizeFields( const MeshBase<MeshType>& );

	void resizeStorage();

	void calculateStorage();

	void initialiseCorners();

	void initialiseImp();

	void verifyMatchingBoundaryConditions( const BoundaryPatch& ) const;

	void checkFieldSizesImp() const;

	void validateImp() const;

	auto calculateTimeStepImp() const -> scalar;

	void calculateFluxesImp();

	void calculateSourcesImp();
	
	void calculateNewStateImp();

	void updateBoundariesImp();

	constexpr static auto zFieldNameImp() -> FieldName {
		return FieldName::BottomElevation;
	}

	template<bool atomicWrite>
	void registerFlux(
		Ref<Array3s> fluxL,
		Ref<Array3s> fluxR,
		const typename MeshBase<MeshType>::EdgeType& edge,
		const Array2i& lr,
		const MeshBase<MeshType>& mesh
	);

	template<bool atomicWrite>
	void registerFlux(
		Ref<Array3s> flux,
		const typename MeshBase<MeshType>::EdgeType& edge,
		const Array2i& lr,
		const MeshBase<MeshType>& mesh
	);

	template<bool secondOrder>
	auto fullyUpwindCT(
		const typename MeshBase<MeshType>::EdgeType& edge,
		Index cell
	) -> const Array3s;

	void slopeSourceToFlux(
		Ref<Array3Ns<2>> flux,
		RowArray2s& zLR,
		const RowArray2s& dLR,
		const RowArray2s& hLR,
		const RowArray2s& zCellLR,
		const RowArray2s& dCellLR,
		const typename MeshBase<MeshType>::EdgeType& edge
	) const;

	template<ReconstructionType recType, EdgeRange orien = EdgeRange::undefined>
	EdgeValues reconstruction(
		const MeshBase<MeshType>& mesh,
		const typename MeshType::EdgeType& edge,
		const Array2i& lr,
		const RowArray2s& zCellLR,
		const RowArray2s& dCellLR
	);

	template<ReconstructionType recType, bool atomicWrite>
	void hllcScheme(const MeshBase<MeshType>& mesh, Index e);

	template<bool atomicWrite>
	void fullyUpwind(const MeshBase<MeshType>& mesh, Index e);

	friend SolverBase<SWESolver<MeshType>>;
	friend swe::SolverThread<MeshType>;
};

} // end namespace hms

#include "sweSolver.tpp"
#include "sweSolver_reconstruction.tpp"
#include "sweSolver_fluxSchemes_edge.tpp"

#endif
